import numpy as np
import matplotlib.pyplot as plt

class Population:
    """ Simulation of a population.

    Attributes:
        popInit: the initial population
        birthrate: the birthrate for 1 year (not timestep!)
        deathrate: the deathrate for 1 year
        timeInit: time at which calculations start
        timeFinal: time at which calculations stop
        timeStep: time steps between 1 year. Should be floating point representation of
                    1 devided by desired sub-steps (e.g. 0.25)
        popList: list that stores the population calculation results for each time step
        timeList: list that stores the times for popList
    """


    
    def __init__(self, popInit, birthrate, deathrate, timeInit, timeFinal, timeStep):

        self.popInit = popInit          # Initial population
        self.pop = popInit              # Set population to inital
        self.birthrate = birthrate      # birthrate
        self.deathrate = deathrate      # deathrate
        
        self.timeInit = timeInit        # initial time
        self.time = timeInit            # time counter
        self.timeFinal = timeFinal      # final time
        self.timeStep = timeStep        # timesteps

        self.popList = list()           # stores results for all time steps
        self.timeList = list()          # stores timesteps for each value


    def setPop(self):
        """ calculates population for one step """
        self.pop = self.pop + self.birth() - self.death()   # new population
        self.time = self.time + self.timeStep               # new time
        self.popList.append(self.pop)       # write new pop value in its list
        self.timeList.append(self.time)     # write new time value in its list    


    def birth(self):
        """ calculates births for current time step """
        return self.pop*self.birthrate*self.timeStep


    def death(self):
        """ calculate deaths for current time step """
        return self.pop * self.deathrate*self.timeStep


    def getPop(self):
        """ print current time and pop """
        print("time:", self.time, "population:", self.pop)


    def getPopAt(self, time):
        """ calculate population for a specific time

        Args:
            time: The time for which the population shall be calculated
        """
        
        self.pop = self.popInit     # reset
        self.time = self.timeInit   # reset

        i = self.timeInit
        while (i < time):     # loop until desired time reached
            self.setPop()
            i = i+self.timeStep

        self.getPop()


    def run(self):
        """ calculates population from initial to final time """
     
        self.pop = self.popInit     # reset
        self.time = self.timeInit   # reset
        self.popList = list()       # reset
        self.timeList = list()      # reset
        self.popList.append(self.pop)  # write initial values in list
        self.timeList.append(self.time) # write initial time in list

        i = self.timeInit
        while (i < self.timeFinal):     # loop until final timestep
            self.setPop()
            self.getPop()
            i = i+self.timeStep


    def plotPop(self):
        """plot the list of calculated populations """

        plot = plt.plot(self.timeList, self.popList)    # create plot from time and pop
        plt.ylabel("Population")                        # set y-axis label
        plt.xlabel("time")                              # set x-axis label
        plt.show(plot)                                  # show plot


